#!/bin/sh

# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# This scripts performs update of stateful partition directories useful for
# dev_mode.

# in order to support developers going from older images with the old shflags
# we'll check for both new shflags and old shflags
if [ -f /usr/share/misc/shflags ] ; then
  . /usr/share/misc/shflags
elif [ -f /usr/lib/shflags ] ; then
  . /usr/lib/shflags
else
  echo >&2 "$0 Unable to source shflags"
  exit 1
fi

# Constants for states.
CLEAN_STATE="clean"
OLD_STATE="old"
RESET_STATE="reset"

LSB_RELEASE="/etc/lsb-release"
STATEFUL_DIR="/mnt/stateful_partition"
UPDATE_STATE_FILE=".update_available"
UPDATE_PROGRESS_FILE="/run/stateful_update/progress"

DEFINE_string stateful_change "${OLD_STATE}" \
  "The state of the new stateful partition - used in update testing."

FLAGS "$@" || exit 1

# Die on error.
set -e

remove_quotes() {
  echo "$1" | sed -e "s/^'//; s/'$//"
}

download_stateful_payload(){
    mkdir -p "$(dirname "$UPDATE_PROGRESS_FILE")"
  # Download and unzip directories onto the stateful partition.
  eval "curl -L --progress-bar --max-time 600 \"${stateful_update_url}\" 2>\"${UPDATE_PROGRESS_FILE}\"" |
      tar --ignore-command-error --overwrite --directory=${STATEFUL_DIR} -xz
}

update_dev_image () {
  local base_update_url devserver_url
  local local_payload_path
  local path=$(remove_quotes "${FLAGS_ARGV}")

  if [ -z "${path}" ]; then
    if [ -f "${STATEFUL_DIR}${LSB_RELEASE}" ]; then
      devserver_url=$(grep CHROMEOS_DEVSERVER ${STATEFUL_DIR}${LSB_RELEASE} |
          cut -f 2 -d '=')
    fi
    if [ -z "${devserver_url}" ]; then
      devserver_url=$(grep CHROMEOS_DEVSERVER ${LSB_RELEASE} | cut -f 2 -d '=')
    fi
    # Sanity check.
    if [ -z "${devserver_url}" ]; then
      echo >&2 "No CHROMEOS_DEVSERVER URL found in lsb-release file."
      exit 1
    fi
    # Devserver URL should never contain "/update"
    devserver_url=$(echo ${devserver_url} | sed -e 's#/update##')
    base_update_url="${devserver_url}/static"
  # Determine whether the path is local.
  elif [ -f "${path}" ] || [ "${path}" = "-" ]; then
    local_payload_path=${path}
  else
    base_update_url=${path}
  fi

  if [ -n "${base_update_url}" ] || [ -n "${FYDEOS_STATEFUL_UPDATE_URL}" ]; then
    local stateful_update_url="${base_update_url}/stateful.tgz"
    if [ -n "$FYDEOS_STATEFUL_UPDATE_URL" ]; then
        stateful_update_url="$FYDEOS_STATEFUL_UPDATE_URL"
    fi
    local download_exit_code=0
    for i in `seq 1 2`; do
      if [ $i -eq 1 ]; then
        echo >&2 "Downloading stateful payload from ${stateful_update_url}"
      else
        echo >&2 "Downloading failed, retrying."
      fi
      if download_stateful_payload; then
        download_exit_code=$?
        echo >&2 "Downloading command returns code ${download_exit_code}."
        break
      else
        download_exit_code=$?
        echo >&2 "Downloading command returns code ${download_exit_code}."
      fi
    done
    if [ ${download_exit_code} -ne 0 ]; then
      return ${download_exit_code}
    fi
    echo >&2 "Successfully downloaded update."
  else
    echo >&2 "Reading local payload ${local_payload_path}"
    # Set timeout to four minutes to avoid waiting on stdin indefinitely.
    timeout 240s tar --ignore-command-error --overwrite \
        --directory=${STATEFUL_DIR} -xzf ${local_payload_path}
    exit_code=$?
    if [ "${exit_code}" -ne 0 ]; then
      echo >&2 "Failed to unzip and write the stateful update."
      return "${exit_code}"
    fi
    echo >&2 "Successfully retrieved update."
  fi

  if [ ! -d "${STATEFUL_DIR}/var_new" ] ||
      [ ! -d "${STATEFUL_DIR}/dev_image_new" ]; then
    echo >&2 "Missing var or dev_image in stateful payload."
    return 1
  fi
}

reset_state () {
  echo >&2 "Resetting stateful update state."
  rm -f "${STATEFUL_DIR}/${UPDATE_STATE_FILE}"
  rm -rf "${STATEFUL_DIR}/var_new"
  rm -rf "${STATEFUL_DIR}/dev_image_new"
}

update_old_state () {
  echo >&2 "Performing standard stateful update."
  echo -n "" > "${STATEFUL_DIR}/${UPDATE_STATE_FILE}"
}

update_clean_state () {
  echo >&2 "Restoring state to factory_install with dev_image."
  echo -n "clobber" > "${STATEFUL_DIR}/${UPDATE_STATE_FILE}"
}

main () {
  if [ "${FLAGS_stateful_change}" = "${RESET_STATE}" ]; then
    reset_state
  elif update_dev_image; then
    if [ "${FLAGS_stateful_change}" = "${OLD_STATE}" ]; then
      update_old_state
    elif [ "${FLAGS_stateful_change}" = "${CLEAN_STATE}" ]; then
      update_clean_state
    else
      echo >&2 "Invalid state given to stateful update.  Aborting..."
      return 1
    fi
  else
    return 1
  fi
}

main $@
